wave_data=data
data_type=shard
gpu_list="0"

export CUDA_VISIBLE_DEVICES="${gpu_list}"
echo "CUDA_VISIBLE_DEVICES is ${CUDA_VISIBLE_DEVICES}"

average_checkpoint=false
average_num=5

stage=5 # start from 0 if you need to start from data preparation
stop_stage=5
dir=exp/final_work_stage6
decode_checkpoint=$dir/epoch_22.pt
recog_set="test/aishell1 test/aishell2 test/test_net test/asru test/test_meeting test/SPEECHIO_ASR_ZH00002 test/SPEECHIO_ASR_ZH00001 test/SPEECHIO_ASR_ZH00000 test/SPEECHIO_ASR_ZH00003 test/SPEECHIO_ASR_ZH00004"

lm_dir=/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_aishell/examples/aishell/do_goutu/data_out

if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
  # Test model, please specify the model you want to test by --checkpoint
  # TODO, Add model average here
  mkdir -p $dir/test
  if [ ${average_checkpoint} == true ]; then
    decode_checkpoint=$dir/avg_${average_num}.pt
    echo "do model average and final checkpoint is $decode_checkpoint"
    python wenet/bin/average_model.py \
      --dst_model $decode_checkpoint \
      --src_path $dir  \
      --num ${average_num} \
      --val_best
  fi
  # Specify decoding_chunk_size if it's a unified dynamic chunk trained model
  # -1 for full chunk
  decoding_chunk_size=
  lm_scale=0.7
  decoder_scale=0.1
  r_decoder_scale=0.7
  ctc_weight=0.5
  decode_modes="hlg_onebest hlg_rescore"
  for test in $recog_set; do
    result_dir=$dir/${test}
    python wenet/bin/recognize.py --gpu $gpu_list \
      --modes $decode_modes \
      --config $dir/train.yaml \
      --data_type $data_type \
      --test_data $wave_data/$test/shard.list \
      --checkpoint $decode_checkpoint \
      --beam_size 10 \
      --batch_size 16 \
      --blank_penalty 0.0 \
      --result_dir $result_dir \
      --ctc_weight $ctc_weight \
      --word $lm_dir/words.txt \
      --hlg $lm_dir/HLG.pt \
      --lm_scale $lm_scale \
      --decoder_scale $decoder_scale \
      --r_decoder_scale $r_decoder_scale \
      ${decoding_chunk_size:+--decoding_chunk_size $decoding_chunk_size}

    for mode in $decode_modes; do
      test_dir=$result_dir/$mode
      python tools/compute-wer.py --char=1 --v=1 \
        $wave_data/$test/text $test_dir/text > $test_dir/wer
    done
  done
fi










